"""
定时创建验收任务
从 clip_annotated_record 中获取未创建的 clip ，按标注类型和批次组织成验收任务
所有的任务都按照车天进行组织，普通任务最多 10 个 clip 一组，
21 对应的任务最多 100 个 clip 一组，并支持批量通过
"""

from collections import defaultdict
from datetime import datetime

from pymongo import UpdateMany, InsertOne, UpdateOne
import pytz
from tqdm import tqdm
from connections import mongo_db, mongo_client

ac_task_col = mongo_db["annotated_acceptance_task"]
check_frame_label_col = mongo_db["check_frame_label"]
clip_anno_col = mongo_db["clip_annotated_record"]


annotate_types = [
    ("parking_free_space", 21),
    ("3D_LANE", 2),
    ("parking_surround_space_detection", 19),
    ("parking_road_sign_mark", 20),
    ("PARKING_3D_LANE", 31),
    ("CITY_4D_LANE", 29),
]


def get_utc_time() -> datetime:
    """
    获取当前UTC时间
    """
    utc_tz = pytz.utc
    now = datetime.now()
    utc_time = now.astimezone(utc_tz)

    return utc_time


def main():
    # 对每个类别的 clip，按照车天进行分组
    for anno_type in annotate_types:
        print(f"\nhandle {anno_type}")
        anno_type_name, anno_type_num = anno_type
        checked_clip_ids = check_frame_label_col.distinct(
            "clip_id", {"annotate_type": anno_type_num}
        )
        ac_tasks = ac_task_col.find(
            {"annotate_type": anno_type_num}, {"clip_ids": 1}
        ).to_list()
        ac_task_clip_ids = set()
        for task in ac_tasks:
            ac_task_clip_ids.update(task["clip_ids"])
        print(
            f"checked clip ids: {len(checked_clip_ids)}, ac task clip ids: {len(ac_task_clip_ids)}"
        )
        diff = set(checked_clip_ids) - ac_task_clip_ids
        print(f"diff clip ids: {len(diff)}")

        # query = {
        #     "annotate_type": anno_type_num,
        #     "state": 2,  # 确认
        #     "accept_status": None,
        # }
        # if anno_type_num != 21:
        #     query["is_vis_processed"] = True
        query = {
            "annotate_type": anno_type_num,
            "clip_id": {"$in": list(diff)},
        }
        # 只获取关注的部分数据
        projection = {
            "clip_id": 1,
            "name": 1,
            "collect_time": 1,
            "plate_no": 1,
            "annotate_type": 1,
            "frame_cnt": 1,
            "car_type": 1,
            "supplier": 1,
        }
        cur_type_clip_anno_list = clip_anno_col.find(query, projection).to_list()
        print(f"find total clip_anno: {len(cur_type_clip_anno_list)}")
        if len(cur_type_clip_anno_list) == 0:
            continue

        exist_ac_task_cnt = ac_task_col.count_documents(
            {"annotate_type": anno_type_num}
        )
        print(f"exist ac task cnt: {exist_ac_task_cnt}")

        # 按照车天进行分组
        vehicle_day_group = defaultdict(list)
        for clip in cur_type_clip_anno_list:
            vehicle_day = str(datetime.fromtimestamp(clip["collect_time"] / 1e6).date())
            key = (clip["plate_no"], vehicle_day)
            vehicle_day_group[key].append(clip)

        # 非 21 的话，10个一组，21 的话，100个一组
        CHUNK = 10 if anno_type_num != 21 else 100
        ac_task_insert_ops = []
        clip_anno_update_ops = []
        task_idx = exist_ac_task_cnt
        for veh_day, clips in vehicle_day_group.items():
            print(f"vehicle day: {veh_day}, clip count: {len(clips)}")
            for ix in range(0, len(clips), CHUNK):
                task_idx += 1
                cur_task_clips = clips[ix : ix + CHUNK]
                new_task = {
                    "task_name": f"[{anno_type_num}][{anno_type_name.upper()}][{veh_day[0]}][{veh_day[1]}]验收任务{task_idx}",
                    "task_description": "旧portal进行的标注，从check_frame_label中重新生成验收任务",
                    "status": 3,  # 旧任务都算做已完成
                    "project": 6,
                    "frequency": "1Hz",
                    "annotate_type": anno_type_num,
                    "is_vis_processed": True,
                    "clip_ids": [clip["clip_id"] for clip in cur_task_clips],
                    "clip_ids_num": len(cur_task_clips),
                    "clip_annotated_records": cur_task_clips,
                    "priority": 5,  # 默认最低优先级为5
                    "created_at": get_utc_time(),
                    "remark": "rebuild from check_frame_label",
                }
                ac_task_insert_ops.append(InsertOne(new_task))
                clip_anno_update_ops.append(
                    UpdateMany(
                        {
                            "clip_id": {"$in": new_task["clip_ids"]},
                            "annotate_type": anno_type_num,
                        },
                        {"$set": {"accept_status": 20}},  # set status pending
                    )
                )
        print("ready to insert task")
        try:
            with mongo_client.start_session() as session, session.start_transaction():
                ret = ac_task_col.bulk_write(
                    ac_task_insert_ops, ordered=False, session=session
                )
                assert ret.inserted_count == len(ac_task_insert_ops)
                ret = clip_anno_col.bulk_write(
                    clip_anno_update_ops, ordered=False, session=session
                )
                assert ret.modified_count == len(cur_type_clip_anno_list)
        except Exception as e:
            print("------------------ exception error", str(e))
        print(f"insert task success, total: {len(ac_task_insert_ops)}")


def main2():
    query = {"remark": {"$exists": True}}
    projection = {"annotate_type": 1, "clip_ids": 1}
    total = ac_task_col.count_documents(query)
    for task in tqdm(ac_task_col.find(query, projection), total=total):
        anno_type, clip_ids = task["annotate_type"], task["clip_ids"]
        clip_anno_col.update_many(
            {"clip_id": {"$in": clip_ids}, "annotate_type": anno_type},
            {"$set": {"accept_status": 20}},
        )


if __name__ == "__main__":
    # main()
    main2()
